from tensorflow.keras import layers, models
import tensorflow.keras.backend as K
from tensorflow.keras.optimizers import SGD
from tensorflow import keras
import tensorflow as tf
from .resnet_tool import conv_block_SE, identity_block_SE, MHSA
from .Tool import compute_output_shape, in_train_phase
from .Estimate import Estimate_BBN
import cv2
import numpy as np


class input_process(layers.Layer):
    def __init__(self, **kwargs):
        super(input_process, self).__init__(**kwargs)

    def call(self, inputs, training=None):
        img_cb = inputs[0]
        img_rb = inputs[1]

        img_train = tf.concat([img_rb, img_rb], axis=0)
        img_val = img_cb

        def train_out():
            return img_train

        def val_out():
            return img_val

        return K.in_train_phase(train_out, val_out, training)

    def compute_output_shape(self, input_shape):
        return compute_output_shape(self, input_shape)


class cb_rb_layer(layers.Layer):
    def __init__(self, **kwargs):
        super(cb_rb_layer, self).__init__(**kwargs)

    def call(self, inputs, training=None):
        def train_out():
            fea_cb, fea_rb = tf.split(inputs, num_or_size_splits=2, axis=0)
            return [fea_cb, fea_rb]

        def val_out():
            return [inputs, inputs]

        return K.in_train_phase(train_out, val_out, training=training)

    def compute_output_shape(self, input_shape):
        return compute_output_shape(self, input_shape)


class BBN(layers.Layer):
    def __init__(self, **kwargs):
        super(BBN, self).__init__(**kwargs)
        # self.feature_exr_fc = feature_extr_fc
        self.alpha = self.add_weight(name="alpha", shape=[], dtype=tf.float32, trainable=False,
                                     initializer="random_normal")
        self.global_avg_pool = layers.GlobalAveragePooling2D()

    def call(self, inputs, training=None):
        cb_fea = inputs[0]
        rb_fea = inputs[1]

        def train_BBN():
            outs = [cb_fea * self.alpha * 2, rb_fea * (1 - self.alpha) * 2]
            out = tf.concat(outs, axis=-1)
            out = self.global_avg_pool(out)
            return out

        def val_BBN():
            out1 = tf.concat([cb_fea, rb_fea], axis=-1)
            out1 = self.global_avg_pool(out1)
            return out1

        return in_train_phase(train_BBN, val_BBN, training)

    def get_config(self):
        config = super(BBN, self).get_config()
        return config

    def compute_output_shape(self, input_shape):
        return compute_output_shape(self, input_shape)
        # return [(n, self.num_class), (n, self.num_class)]

    @property
    def alpha_tensor(self):
        return self.alpha


def BBNBotNet50(input_tensor=None):
    # Determine proper input shape
    if K.image_data_format() == 'channels_last':
        bn_axis = 3
    else:
        bn_axis = 1

    x = layers.ZeroPadding2D(padding=(3, 3), name='conv1_pad')(input_tensor)
    x = layers.Conv2D(64, (7, 7),
                      strides=(2, 2),
                      padding='valid',
                      name='conv1')(x)
    x = layers.BatchNormalization(axis=bn_axis, name='bn_conv1')(x)
    x = layers.Activation('relu')(x)
    x = layers.MaxPooling2D((3, 3), strides=(2, 2))(x)

    x = conv_block_SE(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1))
    x = identity_block_SE(x, 3, [64, 64, 256], stage=2, block='b')
    x = identity_block_SE(x, 3, [64, 64, 256], stage=2, block='c')

    x = conv_block_SE(x, 3, [128, 128, 512], stage=3, block='a')
    x = identity_block_SE(x, 3, [128, 128, 512], stage=3, block='b')
    x = identity_block_SE(x, 3, [128, 128, 512], stage=3, block='c')
    x = identity_block_SE(x, 3, [128, 128, 512], stage=3, block='d')

    x = conv_block_SE(x, 3, [256, 256, 1024], stage=4, block='a', strides=(1, 1))
    x = identity_block_SE(x, 3, [256, 256, 1024], stage=4, block='b')
    x = identity_block_SE(x, 3, [256, 256, 1024], stage=4, block='c')
    x = identity_block_SE(x, 3, [256, 256, 1024], stage=4, block='d')
    x = identity_block_SE(x, 3, [256, 256, 1024], stage=4, block='e')
    x = identity_block_SE(x, 3, [256, 256, 1024], stage=4, block='f')
    return x


def last_block(input, name):
    x = conv_block_SE(input, 3, [512, 512, 2048], stage=5, block='a_' + name, strides=(1, 1))
    x = identity_block_SE(x, 3, [512, 512, 2048], stage=5, block='b_' + name)
    x = identity_block_SE(x, 3, [512, 512, 2048], stage=5, block='c_' + name)
    # model = models.Model(inputs=input,outputs=x)
    return x


def create_model(batch_size, input_size, classes):
    # input_size = args.input_size
    # batch_size = args.batch_size
    input_shape = (input_size, input_size, 3)
    input_img_cb = layers.Input(shape=input_shape)
    input_img_rb = layers.Input(shape=input_shape)
    label_cb = layers.Input(shape=(classes))
    label_rb = layers.Input(shape=(classes))
    weight_cb = layers.Input(shape=())
    weight_rb = layers.Input(shape=())

    input_layer = input_process()
    cb_rb = cb_rb_layer()
    BBN_layer = BBN()

    x = input_layer([input_img_cb, input_img_rb])
    fea = BBNBotNet50(x)
    fea_cb, fea_rb = cb_rb(fea)

    fea_cb = last_block(fea_cb, "cb")
    fea_rb = last_block(fea_rb, "rb")

    x = BBN_layer([fea_cb, fea_rb])
    logic = layers.Dense(units=classes)(x)

    model = models.Model(inputs=[input_img_cb, input_img_rb], outputs=[logic])
    input_map = {"img_cb": input_img_cb, "img_rb": input_img_rb, "label_cb": label_cb, "label_rb": label_rb,
                 "weight_cb": weight_cb, "weight_rb": weight_rb}
    train_output_map = {"acc_cb": tf.reduce_mean(keras.metrics.categorical_accuracy(label_cb, model.output), axis=-1),
                        "acc_rb": tf.reduce_mean(keras.metrics.categorical_accuracy(label_rb, model.output), axis=-1)}
    val_output_map = {"val_acc": tf.reduce_mean(keras.metrics.categorical_accuracy(label_cb, model.output), axis=-1)}
    model = Estimate_BBN(model=model, input_map=input_map, train_output_map=train_output_map,
                         val_output_map=val_output_map)
    model.compile(SGD(0.01), BBN_layer.alpha_tensor)
    return model


def load_model(model_path):
    model = models.load_model(model_path
                              , custom_objects={
            "BBN": BBN,
            "input_process": input_process,
            "cb_rb_layer": cb_rb_layer,
            "tf": tf
        }
                              )
    model.summary()
    return model


def get_predict(model_path, img_size):
    model = load_model(model_path)
    predict_iter = K.function(inputs=model.inputs+[K.learning_phase()], outputs=model.outputs)

    def predict(img_path):
        img = cv2.imread(img_path)
        img = cv2.resize(img, (img_size, img_size))
        img = img / 255.
        res = predict_iter([np.array([img]),np.array([img]),0.])
        return res

    return predict


if __name__ == '__main__':
    model = load_model("D://programing//AILib_tf//out//temp_model-val_acc_0.625.h5")
    model.summary()
